package com.sparrow.common.lock;

import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * 分段读写锁
 */
public class StripedReadWriteLock {

    private final ReentrantReadWriteLock[] locks;
    private int mask = 0x7fffffff;

    public StripedReadWriteLock(int stripes) {
        locks = new ReentrantReadWriteLock[stripes];
        for (int i = 0; i < stripes; i++) {
            locks[i] = new ReentrantReadWriteLock();
        }
    }

    private ReentrantReadWriteLock getLock(Object key) {
        int index = (key.hashCode() & mask) % locks.length;
        return locks[index];
    }

    public void readLock(Object key, Runnable action) {
        ReentrantReadWriteLock lock = getLock(key);
        lock.readLock().lock();
        try {
            action.run();
        } finally {
            lock.readLock().unlock();
        }
    }

    public void writeLock(Object key, Runnable action) {
        ReentrantReadWriteLock lock = getLock(key);
        lock.writeLock().lock();
        try {
            action.run();
        } finally {
            lock.writeLock().unlock();
        }
    }

    // 根据 PK 获取锁段
    private ReentrantReadWriteLock lockFor(Object pk) {
        int index = (pk.hashCode() & mask) % locks.length;
        return locks[index];
    }
}
